# 导入依赖库
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import math


def get_positional_encoding(max_seq_len, embed_dim):
    # pos 如果为0的话就把那一行全部置为0，不用进行计算
    # 第一行全为0
    positional_encoding_h = np.array(
        [
            [pos / np.power(10000, 2 * i / embed_dim) for i in range(embed_dim)]
            if pos != 0 else np.zeros(embed_dim) for pos in range(max_seq_len)
        ])
    # 每两步前进 ：：2
    positional_encoding_h[1:, 0::2] = np.sin(positional_encoding_h[1:, 0::2])  # dim 2i 偶数
    positional_encoding_h[1:, 1::2] = np.cos(positional_encoding_h[1:, 1::2])  # dim 2i+1 奇数
    # 归一化, 用位置嵌入的每一行除以它的模长
    # denominator = np.sqrt(np.sum(position_enc**2, axis=1, keepdims=True))
    # position_enc = position_enc / (denominator + 1e-8)
    return positional_encoding_h


positional_encoding = get_positional_encoding(max_seq_len=100, embed_dim=16)
plt.figure(figsize=(10, 10))
sns.heatmap(positional_encoding)
plt.title("Sinusoidal Function")
plt.xlabel("hidden dimension")
plt.ylabel("sequence length")
plt.show()
